#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import torch
from torch import nn
import torch.nn.functional as F
from abc import ABC
import math


class MLP(nn.Module):
    def __init__(self, dim_in, dim_hidden, dim_out):
        super(MLP, self).__init__()
        self.layer_input = nn.Linear(dim_in, dim_hidden)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout()
        self.layer_hidden = nn.Linear(dim_hidden, dim_out)

    def forward(self, x):
        x = x.view(-1, x.shape[1] * x.shape[-2] * x.shape[-1])
        x = self.layer_input(x)
        x = self.dropout(x)
        x = self.relu(x)
        x = self.layer_hidden(x)
        return x


class CNNMnist(nn.Module):
    def __init__(self, args):
        super(CNNMnist, self).__init__()
        self.conv1 = nn.Conv2d(args.num_channels, 20, kernel_size=5)
        self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(800, 50)
        self.fc2 = nn.Linear(50, args.num_classes)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        #x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
        x = F.relu(self.fc1(x))
        #x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return x


class CNNFashionMnist(nn.Module):
    def __init__(self, args):
        super(CNNFashionMnist, self).__init__()
        self.conv1 = nn.Conv2d(args.num_channels, 16, kernel_size=5)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(512, 256)
        self.fc2 = nn.Linear(256, args.num_classes)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        #x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
        x = F.relu(self.fc1(x))
        #x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return x


class CNNCifar(nn.Module):
    def __init__(self, args):
        super(CNNCifar, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, args.num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class CNNCifar100(nn.Module):
    def __init__(self, args):
        super(CNNCifar100, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, args.num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


class LeNet5Mnist(nn.Module, ABC):
    def __init__(self, args):
        super(LeNet5Mnist, self).__init__()  # super用法:继承父类nn.Model的属性，并用父类的方法初始化这些属性

        self.conv1 = nn.Sequential(
            # nn.Conv2d(in_dim, 6, 5, 1, 2),  # out_dim=6, kernel_size=5, stride=1, padding=2
            nn.Conv2d(args.num_channels, 6, 5, 1, 0),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)  # kernel_size=2, padding=2
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5, 1, 0),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(256, 120),  # in_features=400, out_features=120
            nn.Linear(120, 84),
            nn.Linear(84, args.num_classes)
        )

    def forward(self, x):
        out_conv1 = self.conv1(x)
        out_conv2 = self.conv2(out_conv1)
        out_conv = out_conv2.view(out_conv2.size(0), -1)

        out = self.fc(out_conv)
        return out


class LeNet5Cifar(nn.Module, ABC):
    def __init__(self, args):
        super(LeNet5Cifar, self).__init__()  # super用法:继承父类nn.Model的属性，并用父类的方法初始化这些属性

        self.conv1 = nn.Sequential(
            # nn.Conv2d(in_dim, 6, 5, 1, 2),  # out_dim=6, kernel_size=5, stride=1, padding=2
            nn.Conv2d(args.num_channels, 6, 5, 1, 0),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)  # kernel_size=2, padding=2
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, 5, 1, 0),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(400, 120),  # in_features=400, out_features=120
            nn.Linear(120, 84),
            nn.Linear(84, args.num_classes)
        )

    def forward(self, x):
        out_conv1 = self.conv1(x)
        out_conv2 = self.conv2(out_conv1)
        out_conv = out_conv2.view(out_conv2.size(0), -1)

        out = self.fc(out_conv)
        return out


def lenet5():
    """ return a LeNet 5 object
    """
    return LeNet5Cifar(3, 10)


class ResNet_cifar(nn.Module):
    def __init__(self, args, block, num_blocks, num_classes=10):
        super(ResNet_cifar, self).__init__()
        self.inplanes = 16
        print('num_classes:', num_classes)
        # declare lambda array
        num_block = 0
        for i in num_blocks:
            num_block += i
        print('block number:', num_block)

        self.conv1 = nn.Conv2d(args.num_channels, 16, kernel_size=3, padding=1, bias=False)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)

        self.bn = nn.BatchNorm2d(64 * block.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.avgpool = nn.AvgPool2d(8)
        self.fc = nn.Linear(64 * block.expansion, num_classes)

        for k, m in self.named_modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(0.5)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, num_blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, num_blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        result = {'representation' : x}
        x = self.fc(x)
        result['output'] = x
        return result

class ResNet_mnist(nn.Module):
    def __init__(self, args, block, num_blocks, num_classes=10):
        super(ResNet_mnist, self).__init__()
        self.inplanes = 16
        print('num_classes:', num_classes)
        # declare lambda array
        num_block = 0
        for i in num_blocks:
            num_block += i
        print('block number:', num_block)

        self.conv1 = nn.Conv2d(args.num_channels, 16, kernel_size=3, padding=1, bias=False)
        self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)

        self.bn = nn.BatchNorm2d(64 * block.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size = 3, stride = 1, padding =1)
        self.fc = nn.Linear(64 * block.expansion, num_classes)

        for k, m in self.named_modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(0.5)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, num_blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, num_blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.maxpool(self.conv1(x))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)

        x = F.avg_pool2d(x, 7)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x



class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.bn1 = nn.BatchNorm2d(inplanes)
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes)
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x
        if self.downsample is not None:
            residual = self.downsample(x)

        out = self.bn1(x)
        out = self.relu(out)
        out = self.conv1(out)

        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv2(out)

        out = self.bn3(out)
        out = self.relu(out)
        out = self.conv3(out)

        out += residual
        return out


class BaseBlock(nn.Module):
    alpha = 1

    def __init__(self, input_channel, output_channel, t=6, downsample=False):
        """
            t:  expansion factor, t*input_channel is channel of expansion layer
            alpha:  width multiplier, to get thinner models
            rho:    resolution multiplier, to get reduced representation
        """
        super(BaseBlock, self).__init__()
        self.stride = 2 if downsample else 1
        self.downsample = downsample
        self.shortcut = (not downsample) and (input_channel == output_channel)

        # apply alpha
        input_channel = int(self.alpha * input_channel)
        output_channel = int(self.alpha * output_channel)

        # for main path:
        c = t * input_channel
        # 1x1   point wise conv
        self.conv1 = nn.Conv2d(input_channel, c, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(c)
        # 3x3   depth wise conv
        self.conv2 = nn.Conv2d(c, c, kernel_size=3, stride=self.stride, padding=1, groups=c, bias=False)
        self.bn2 = nn.BatchNorm2d(c)
        # 1x1   point wise conv
        self.conv3 = nn.Conv2d(c, output_channel, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(output_channel)

    def forward(self, inputs):
        # main path
        x = F.relu6(self.bn1(self.conv1(inputs)), inplace=True)
        x = F.relu6(self.bn2(self.conv2(x)), inplace=True)
        x = self.bn3(self.conv3(x))

        # shortcut path
        x = x + inputs if self.shortcut else x

        return x


def ResNet8_cifar(args):
    num_block = [1, 1, 1]
    return ResNet_cifar(args, BaseBlock, num_block, num_classes=args.num_classes)


def ResNet20_cifar(args):
    num_block = [2, 2, 2]
    return ResNet_cifar(args, Bottleneck, num_block, num_classes=args.num_classes)

def ResNet20_mnist(args):
    num_block = [2, 2, 2]
    return ResNet_mnist(args, Bottleneck, num_block, num_classes=args.num_classes)


def ResNet56_cifar(args):
    return ResNet_cifar(args,Bottleneck, [6, 6, 6], num_classes=args.num_classes)


def ResNet101_cifar(args):
    return ResNet_cifar(args,Bottleneck, [11, 11, 11], num_classes=args.num_classes)


def ResNet164_cifar(args):
    return ResNet_cifar(args,Bottleneck, [18, 18, 18], num_classes=args.num_classes)



class bottleneck(nn.Module):
    def __init__(self, inchannels, outchannels, stride, expansion):
        super().__init__()
        self.inchannels = inchannels
        self.outchannels = outchannels
        self.stride = stride

        self.residual = nn.Sequential(
            nn.Conv2d(in_channels=inchannels, out_channels=expansion*inchannels, kernel_size=1),
            nn.BatchNorm2d(expansion*inchannels),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels=inchannels*expansion, out_channels=inchannels*expansion, kernel_size=3, padding=1,
                      groups=inchannels*expansion, stride=stride),
            nn.BatchNorm2d(expansion*inchannels),
            nn.ReLU6(inplace=True),

            nn.Conv2d(in_channels=expansion * inchannels, out_channels=outchannels, kernel_size=1),
            nn.BatchNorm2d(outchannels)
        )

    def forward(self, x):
        out = self.residual(x)

        if self.inchannels == self.outchannels and self.stride == 1:
            out += x
        return out





class MobileNetV2(nn.Module):
    def __init__(self, args, alpha=1):
        super().__init__()

        self.Conv1 = nn.Sequential(
            nn.Conv2d(in_channels=args.num_channels, out_channels=int(alpha*32), kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True)
        )

        self.stage1 = bottleneck(int(alpha*32), 16, 1, 1)
        self.stage2 = self.make_layer(int(alpha*16), 6, int(alpha*24), 2, 2)
        self.stage3 = self.make_layer(int(alpha*24), 6, int(alpha*32), 3, 2)
        self.stage4 = self.make_layer(int(alpha*32), 6, int(alpha*64), 4, 2)
        self.stage5 = self.make_layer(int(alpha*64), 6, int(alpha*96), 3, 1)
        self.stage6 = self.make_layer(int(alpha*96), 6, int(alpha*160), 3, 1)
        self.stage7 = self.make_layer(int(alpha*160), 6, int(alpha*320), 1, 1)

        self.Conv2 = nn.Sequential(
            nn.Conv2d(in_channels=int(alpha*320), out_channels=1280, kernel_size=1),
            nn.BatchNorm2d(1280),
            nn.ReLU6(inplace=True)
        )

        self.AvgPool = nn.AdaptiveAvgPool2d(1)
        self.drop = nn.Dropout(0.2)
        self.Conv3 = nn.Conv2d(in_channels=1280, out_channels=args.num_classes, kernel_size=1)

    def make_layer(self, inchannels, t, outchannels, n, s):
        layer = []
        layer.append(bottleneck(inchannels, outchannels, s, t))
        n = n - 1
        while n:
            layer.append(bottleneck(outchannels, outchannels, 1, t))
            n -= 1
        return nn.Sequential(*layer)

    def forward(self, x):
        out = self.Conv1(x)
        out = self.stage1(out)
        out = self.stage2(out)
        out = self.stage3(out)
        out = self.stage4(out)
        out = self.stage5(out)
        out = self.stage6(out)
        out = self.stage7(out)
        out = self.Conv2(out)
        out = self.AvgPool(out)
        out = self.drop(out)
        out = self.Conv3(out)
        out = out.view(out.size(0), -1)

        return out


def mobilenetv2(args, alpha=1):
    return MobileNetV2(args, 1)


class VGG16(nn.Module):
    def __init__(self, args):
        super(VGG16, self).__init__()
        self.features = nn.Sequential(
            # 1
            nn.Conv2d(args.num_channels, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            # 2
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 3
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            # 4
            nn.Conv2d(128, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 5
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            # 6
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            # 7
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 8
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            # 9
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            # 10
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            # 11
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            # 12
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            # 13
            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.AvgPool2d(kernel_size=1, stride=1),
        )
        self.classifier = nn.Sequential(
            # 14
            nn.Linear(512, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            # 15
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            # 16
            nn.Linear(4096, args.num_classes),
        )
        # self.classifier = nn.Linear(512, 10)

    def forward(self, x):
        out = self.features(x)
        # print('out.shape:',out.shape)
        out = out.view(out.size(0), -1)
        #        print(out.shape)
        out = self.classifier(out)
        #        print(out.shape)
        return out



class VGG16_mnist(nn.Module):
    def __init__(self, args):
        super(VGG16_mnist, self).__init__()
        self.features = nn.Sequential(nn.Conv2d(1, 64, kernel_size=(1, 1), stride=(1, 1), padding=(1, 1)),
                     nn.ReLU(inplace=True),
                     nn.Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
                     nn.ReLU(inplace=True),
                     nn.MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False),
                     nn.Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
                     nn.ReLU(inplace=True),
                     nn.MaxPool2d(kernel_size=2, stride=1, padding=0, dilation=1, ceil_mode=False),
                     nn.Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
                     nn.ReLU(inplace=True),
                     nn.Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
                     nn.ReLU(inplace=True),
                     nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False),
                     nn.AdaptiveAvgPool2d(output_size=(7, 7)),
                     Flatten(),
                     nn.Linear(in_features=25088, out_features=256, bias=True),
                     nn.Linear(in_features=256, out_features=args.num_classes, bias=True)
                     )
        # self.classifier = nn.Linear(512, 10)

    def forward(self, x):
        out = self.features(x)
        # print('out.shape:',out.shape)
        # out = out.view(out.size(0), -1)
        # #        print(out.shape)
        # out = self.classifier(out)
        # #        print(out.shape)
        return out


class Flatten(nn.Module):
    def __init__(self):
        super(Flatten,self).__init__()
    def forward(self,x):
        shape = torch.prod(torch.tensor(x.shape[1:])).item()
        return x.reshape(-1,shape)


# python main_fed.py --dataset cifar --num_channels 3 --model vgg --noniid_case 3 --epochs 600 --gpu 0 --multi_plot_case 6 --num_classes 10 --loss_type cemin05 --iterations 1
# python main_fed.py --dataset cifar --num_channels 3 --model vgg --noniid_case 3 --epochs 600 --gpu 0 --multi_plot_case 6 --num_classes 10 --loss_type cemin05 --iterations 1